import pyaudio
import wave
import numpy as np

RESPEAKER_RATE = 16000
RESPEAKER_CHANNELS = 2
RESPEAKER_WIDTH = 2
# run getDeviceInfo.py to get index
RESPEAKER_INDEX = 2  # refer to input device id
CHUNK = 1024
RECORD_SECONDS = 4
WAVE_OUTPUT_FILENAME = "output_one_channel.wav"

class VoiceRecord:
    def __init__(self):
        self.p = pyaudio.PyAudio()
    
    def rec(self):
        stream = self.p.open(
                    rate=RESPEAKER_RATE,
                    format=self.p.get_format_from_width(RESPEAKER_WIDTH),
                    channels=RESPEAKER_CHANNELS,
                    input=True,
                    input_device_index=RESPEAKER_INDEX,)

        print("* recording")

        frames = [] 

        for i in range(0, int(RESPEAKER_RATE / CHUNK * RECORD_SECONDS)):
            data = stream.read(CHUNK)
            # extract channel 0 data from 2 channels, if you want to extract channel 1, please change to [1::2]
            a = np.frombuffer(data,dtype=np.int16)[1::2]
            frames.append(a.tobytes())

        print("* done recording")

        stream.stop_stream()
        stream.close()
        self.p.terminate()

        wf = wave.open(WAVE_OUTPUT_FILENAME, 'wb')
        wf.setnchannels(1)
        wf.setsampwidth(self.p.get_sample_size(self.p.get_format_from_width(RESPEAKER_WIDTH)))
        wf.setframerate(RESPEAKER_RATE)
        wf.writeframes(b''.join(frames))
        wf.close()


